import numpy as np
from PIL import Image
from skimage.metrics import structural_similarity
import cv2
import os
import multiprocessing
import math

# 用于直方图比较的 Normalize 函数
def normalize(data):
    return data / np.sum(data)

# 使用相关性的直方图相似度（接近 1 的值更相似）
def hist_similarity(img1, img2, hist_size=256):
    imghistb1 = cv2.calcHist([img1], [0], None, [hist_size], [0, 256])
    imghistg1 = cv2.calcHist([img1], [1], None, [hist_size], [0, 256])
    imghistr1 = cv2.calcHist([img1], [2], None, [hist_size], [0, 256])

    imghistb2 = cv2.calcHist([img2], [0], None, [hist_size], [0, 256])
    imghistg2 = cv2.calcHist([img2], [1], None, [hist_size], [0, 256])
    imghistr2 = cv2.calcHist([img2], [2], None, [hist_size], [0, 256])

    distanceb = cv2.compareHist(normalize(imghistb1), normalize(imghistb2), cv2.HISTCMP_CORREL)
    distanceg = cv2.compareHist(normalize(imghistg1), normalize(imghistg2), cv2.HISTCMP_CORREL)
    distancer = cv2.compareHist(normalize(imghistr1), normalize(imghistr2), cv2.HISTCMP_CORREL)
    return np.mean([distanceb, distanceg, distancer])

# 在一个进程中处理一部分 comparison_folders
def process_folder_subset(img_nation_folder, folder_subset, output_queue):
    partial_results = {}

    for nation_image_name in os.listdir(img_nation_folder):
        nation_image_path = os.path.join(img_nation_folder, nation_image_name)
        nation_image = cv2.imread(nation_image_path)

        max_similarity = -1
        most_similar_image = None
        most_similar_folder = None

        for folder in folder_subset:
            for comparison_image_name in os.listdir(folder):
                comparison_image_path = os.path.join(folder, comparison_image_name)
                comparison_image = cv2.imread(comparison_image_path)

                if comparison_image is None or nation_image is None:
                    continue

                # 计算相似度
                similarity = hist_similarity(nation_image, comparison_image)

                # 更新最大相似度
                if similarity > max_similarity:
                    max_similarity = similarity
                    most_similar_image = comparison_image_name
                    most_similar_folder = folder

        # 存储当前图像的结果
        partial_results[nation_image_name] = (most_similar_folder, most_similar_image, max_similarity)

    output_queue.put(partial_results)

# 主函数，使用多进程实现
def find_most_similar_images_multiprocess(img_nation_folder, comparison_folders, num_processes=8):
    # 分割 comparison_folders
    chunk_size = math.ceil(len(comparison_folders) / num_processes)
    folder_chunks = [comparison_folders[i:i + chunk_size] for i in range(0, len(comparison_folders), chunk_size)]

    # 创建进程和队列
    output_queue = multiprocessing.Queue()
    processes = []

    for folder_subset in folder_chunks:
        p = multiprocessing.Process(target=process_folder_subset, args=(img_nation_folder, folder_subset, output_queue))
        processes.append(p)
        p.start()

    # 收集所有进程的结果
    results = {}
    for _ in processes:
        results.update(output_queue.get())

    # 等待所有进程完成
    for p in processes:
        p.join()

    return results

if __name__ == '__main__':
    img_test = './neu_seg_competition/images/test'
    img_training = './neu_seg_competition/images/training'
    img_nation = './neu_seg_competition_nation/Img'

    # 获取比较结果
    comparison_folders = [img_test, img_training]
    results = find_most_similar_images_multiprocess(img_nation, comparison_folders)

    for nation_image, (folder, similar_image, similarity) in results.items():
        print(f"图像 '{nation_image}' 在 img_nation 中最相似的图像是 '{similar_image}' 在文件夹 '{folder}' 中，相似度分数为 {similarity:.4f}")
